import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from sklearn.datasets import fetch_openml
from sklearn.utils import shuffle

# 加载 MNIST 数据集
mnist = fetch_openml("mnist_784", version=1, as_frame=False)
X, y = mnist.data / 255.0, mnist.target.astype(int)

# 打乱数据，并选择子集进行可视化
X, y = shuffle(X, y, random_state=42)
X_train, y_train = X[:5000], y[:5000]  # 子集选择

# PCA 降维到 3 维
pca = PCA(n_components=3, svd_solver='randomized', random_state=42)
X_transformed = pca.fit_transform(X_train)

# 可视化前两维数据
fig, ax = plt.subplots(figsize=(10, 10))  # 设置合理的图像大小
scatter = ax.scatter(X_transformed[:, 0], X_transformed[:, 1], c=y_train, cmap='prism', alpha=0.6)
ax.set_title("MNIST Data Visualization with PCA")
ax.set_xlabel("Principal Component 1")
ax.set_ylabel("Principal Component 2")

# 添加颜色条
colorbar = fig.colorbar(scatter, ax=ax, orientation='vertical')
colorbar.set_label('Digit Label')

# 保存图片
plt.savefig("images/mnist_pca.png", dpi=300)
plt.show()
